#include <stdio.h>
#include <assert.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <mpi.h>

#include "utils.hpp"
#include "precond.hpp"
#include "iter_solver.hpp"
#include "my_parcsr.hpp"

#define MIN(a, b) (a) < (b) ? (a) : (b)

// 测试循环次数
#ifndef TEST_CNT
#define TEST_CNT 1
#endif

// 数据路径
const char * dir_name = "/home/2023-fall/2023310751/data"; // TODO: 在此处修改数据路径

// 记录测试信息
typedef struct {
  double setup, solve, prec;
  int iter;
} TEST_RECORD;

int main(int argc, char * argv[])
{
    int my_pid, num_procs, cnt = 1;

    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &num_procs);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);

    // case_name = [kind]/[var]_101; kind = kvlcc2, suboff; var = u, v, w
    // e.g. "kvlcc2/u_101" 是一个合法的 case_name
    const char * case_name = argv[cnt++]; 
    const int glb_nrows = atoi(argv[cnt++]); // 矩阵的全局行数
    const char * prec_name = argv[cnt++]; // 预条件子名称，可选 Jacobi, PGS
    double tolerance = 1.0e-5; // 目标残差与初始残差的比值
    if (argc > 4) tolerance = atof(argv[cnt++]);

    // 校验 case_name
    if (!(strstr(case_name, "kvlcc2") != NULL || strstr(case_name, "suboff") != NULL)) {
        if (my_pid == 0) printf("Error: Unrecognized case name %s\n", case_name);
        MPI_Abort(MPI_COMM_WORLD, -1);
    }
    
    // 分配本进程负责的行号 [ilower, iupper]
    int loc_nrows = glb_nrows / num_procs;
    int ilower = my_pid * loc_nrows;
    if (glb_nrows > loc_nrows * num_procs) {
        int remain_nblks = glb_nrows - loc_nrows * num_procs;
        if (my_pid < remain_nblks) {
            loc_nrows++;
        }
        ilower += MIN(my_pid, remain_nblks);
    }
    int iupper = ilower + loc_nrows - 1;

    // 读入数据
    int ret;
    char pathname[200], filename[200];
    sprintf(pathname, "%s/%s", dir_name, case_name); 
    if (my_pid == 0) printf("reading data from %s\n", pathname);

    // 读入二进制的矩阵数据
    // row_ptr: Ai.bin
    // col_idx: Aj.bin
    // vals: Av.bin

    int * dist_row_ptr = (int *) malloc(sizeof(int) * (loc_nrows+1));// 分布式存储的行指针（数值为全局号）
    sprintf(filename, "%s/Ai.bin", pathname);
    if ((ret = read_binary(dist_row_ptr, filename, sizeof(int), ilower, loc_nrows+1)) != loc_nrows+1) {
        printf("Error! not enough rows\n"); // 错误处理
        MPI_Abort(MPI_COMM_WORLD, 1);
    }
    const int loc_nnz = dist_row_ptr[loc_nrows] - dist_row_ptr[0];// 本进程负责的行内一共的非零元数
    int * dist_col_idx = (int *) malloc((loc_nnz) * sizeof(int));// 分布式存储的列序号（数值为全局号）
    double* dist_vals = (double*)malloc(loc_nnz * sizeof(double));
    sprintf(filename, "%s/Aj.bin", pathname);
    if ((ret = read_binary(dist_col_idx, filename, sizeof(int), dist_row_ptr[0], loc_nnz)) != loc_nnz) {
        printf("Error! not enough dist_col_idx: %d\n", ret);
        MPI_Abort(MPI_COMM_WORLD, 2);
    }
    sprintf(filename, "%s/Av.bin", pathname);
    if ((ret = read_binary(dist_vals, filename, sizeof(double), dist_row_ptr[0], loc_nnz)) != loc_nnz) {
        printf("Error! not enough dist_vals: %d\n", ret);
        MPI_Abort(MPI_COMM_WORLD, 3);
    }
    check_input(dist_vals   , loc_nnz); // 检查输入是否合法

    // 读入二进制的向量数据
    // b: b.bin
    // x0: x0.bin

    double* dist_b = (double*) malloc (loc_nrows * sizeof(double));
    double* dist_x = (double*) malloc (loc_nrows * sizeof(double));
    sprintf(filename, "%s/b.bin", pathname);
    if ((ret = read_binary(dist_b, filename, sizeof(double), ilower, loc_nrows)) != loc_nrows) {
        printf("Error! not enough b\n");
        MPI_Abort(MPI_COMM_WORLD, 4);
    }
    sprintf(filename, "%s/x0.bin", pathname);
    if ((ret = read_binary(dist_x, filename, sizeof(double), ilower, loc_nrows)) != loc_nrows) {
        if (my_pid == 0) printf(" ====> use zero inital guess instead!\n");
        for (int i = 0; i < loc_nrows; i++)
            dist_x[i] = 0.0;
    }
    check_input(dist_b, loc_nrows);
    check_input(dist_x, loc_nrows);

    // 校验：向量点积
    {
        double self_dot[2] = {0.0, 0.0};
        for (int i = 0; i < loc_nrows; i++) {
            self_dot[0] += (double) dist_b[i] * (double) dist_b[i];
            self_dot[1] += (double) dist_x[i] * (double) dist_x[i];
        }
        double glb_dot[2];
        MPI_Allreduce(self_dot, glb_dot, 2, MPI_DOUBLE, MPI_SUM, MPI_COMM_WORLD);
        if (my_pid == 0) {
            printf("My calc dot\n");
            printf("(  b,   b) = %.15e\n",  glb_dot[0]);
            printf("(  x,   x) = %.15e\n",  glb_dot[1]);
        }
    }

    // 初始化 ParCSR 矩阵和向量
    MyParCSRMatrix<int, double, double> parcsr_A(MPI_COMM_WORLD, glb_nrows, ilower, iupper + 1,
        dist_row_ptr, dist_col_idx, dist_vals);
    MyParCSRVector<int, double> parcsr_b(MPI_COMM_WORLD, glb_nrows, ilower, iupper + 1, dist_b);
    MyParCSRVector<int, double> parcsr_x(MPI_COMM_WORLD, glb_nrows, ilower, iupper + 1, dist_x);
    MyParCSRVector<int, double> parcsr_y(MPI_COMM_WORLD, glb_nrows, ilower, iupper + 1, nullptr);

    // 计算基础校验值
    double b_dot = vec_dot(parcsr_b, parcsr_b);
    double x_dot = vec_dot(parcsr_x, parcsr_x);
    parcsr_A.Mult(1.0, parcsr_b, 0.0, parcsr_y);;
    double Ab_dot = vec_dot(parcsr_y, parcsr_y);
    parcsr_A.Mult(1.0, parcsr_x, 0.0, parcsr_y);
    double Ax_dot = vec_dot(parcsr_y, parcsr_y);
    parcsr_A.Mult(-1.0, parcsr_x, 1.0, parcsr_b, parcsr_y);
    double r0_norm = vec_norm(parcsr_y);

    if (my_pid == 0) {
        printf("(  b,   b) = %.20e\n",  b_dot);
        printf("(  x,   x) = %.20e\n",  x_dot);
        printf("(A*b, A*b) = %.20e\n", Ab_dot);
        printf("(A*x, A*x) = %.20e\n", Ax_dot);
        printf("||r0||_2 = %.15e\n", r0_norm);
    }

    /******************************************* 数据读取结束 ****************************************/

    TEST_RECORD records[TEST_CNT];
    double t_setup, t_calc, t_pre;
    int num_iterations;
    double final_res_norm;

    // 测试循环
    for (int test = 0; test < TEST_CNT; test++) {
        // 拷贝初始解
        MyParCSRVector<int, double> parcsr_sol(MPI_COMM_WORLD, glb_nrows, ilower, iupper + 1, dist_x);

        if (strstr(case_name, "kvlcc2") || strstr(case_name, "suboff")) {
            Precond<int, pc_t, ksp_t, ksp_t> * precond = nullptr;
            IterSolver<int, pc_t, ksp_t, ksp_t> * solver = nullptr;

            // 初始化预条件子
            if (strstr(prec_name, "Jacobi")) {
                precond = new Jacobi<int, pc_t, ksp_t, ksp_t>();
            } else if (strstr(prec_name, "PGS")) {
                precond = new PGS<int, pc_t, ksp_t, ksp_t>();
            }
            assert(precond);
            
            // 初始化外迭代
            solver = new GMRESSolver<int, pc_t, ksp_t, ksp_t>();
            solver->SetPrecond(precond);
            solver->SetMaxIter(50);
            solver->SetAbsTol(r0_norm * tolerance);
            solver->SetRelTol(0.0);
            ((GMRESSolver<int, pc_t, ksp_t, ksp_t> *)solver)->SetRestartlen(10);

            // 根据矩阵信息建立预条件子
            MPI_Barrier(MPI_COMM_WORLD);
            t_setup = MPI_Wtime();
            solver->Setup(parcsr_A);
            t_setup = MPI_Wtime() - t_setup;

            // 迭代求解
            MPI_Barrier(MPI_COMM_WORLD);
            t_calc = MPI_Wtime();
            solver->Solve(parcsr_b, parcsr_sol);
            t_calc = MPI_Wtime() - t_calc;

            // 验证正确性(计算真实残差)
            parcsr_A.Mult(-1.0, parcsr_sol, 1.0, parcsr_b, parcsr_y);
            double r_norm = vec_norm(parcsr_y);
            if (my_pid == 0) {
                printf("||b - A * x_0|| / ||b|| = %.20e\n", r0_norm / sqrt(b_dot));
                printf("||b - A * x  || / ||b|| = %.20e\n", r_norm / sqrt(b_dot));
            }

            // 运行信息
            num_iterations = solver->GetNumIterations();
            final_res_norm = solver->GetFinalNorm();
            if (my_pid == 0) {
                printf("\n");
                printf("Iterations = %d\n", num_iterations);
                printf("Prepare cost %.5f\n", t_pre);
                printf("Time cost %.5f %.5f %.5f %d\n", t_setup, t_calc, t_setup + t_calc, num_iterations);
                printf("Spmv: %.5f, Prec: %.5f\n", solver->GetSpmvTime(), solver->GetPrecTime());
                printf("Final Relative Residual Norm = %e\n", final_res_norm);
                printf("\n");
            }

            delete precond;
            delete solver;

            if (test == TEST_CNT - 1) {
                MPI_Datatype etype = (sizeof(ksp_t) == 4) ? MPI_FLOAT : MPI_DOUBLE;
                ksp_t *sol = nullptr;
                int *recvcount = nullptr, *displs = nullptr;
                if (my_pid == 0) {
                    sol = (ksp_t *) malloc(sizeof(ksp_t) * glb_nrows);
                    recvcount = (int *) malloc(sizeof(int) * num_procs);
                    displs = (int *) malloc(sizeof(int) * num_procs);
                    for (int i = 0; i < num_procs; i++) {
                        int loc_nrows_i = glb_nrows / num_procs;
                        int ilower_i = i * loc_nrows_i;
                        if (glb_nrows > loc_nrows_i * num_procs) {
                            int remain_nblks = glb_nrows - loc_nrows_i * num_procs;
                            if (i < remain_nblks) {
                                loc_nrows_i++;
                            }
                            ilower_i += MIN(i, remain_nblks);
                        }
                        recvcount[i] = loc_nrows_i;
                        displs[i] = ilower_i;
                    }
                }
                MPI_Gatherv(parcsr_sol.data, loc_nrows, etype, sol, recvcount, displs, etype, 0, MPI_COMM_WORLD);
                if (my_pid == 0) {
                    sprintf(filename, "%s/ans.dat", pathname);
                    FILE *fp = fopen(filename, "w");
                    if (fp == NULL) {
                        printf("Failed to open file.\n");
                        exit(1);
                    }
                    for (int i = 0; i < glb_nrows; i++)
                        fprintf(fp, "%.27e\n", sol[i]);
                    fclose(fp);
                }
            }
        }
        else assert(false);
        
        records[test].iter = num_iterations;
        records[test].setup = t_setup;
        records[test].solve = t_calc;
        
    }// test loop
    
    // 打印运行时间
    // 10次中 最好(BEST)、平均(Average)、首次(First) 的运行时间
    if (my_pid == 0) {
        int best = 0;
        double best_tot = records[0].setup + records[0].solve;
        double avg_setup = 0, avg_solve = 0, avg_tot = 0;
        for (int i = 0; i < TEST_CNT; i++) {
            if (best_tot > (records[i].setup + records[i].solve)) {
                best_tot =  records[i].setup + records[i].solve;
                best = i;
            }
            avg_setup += records[i].setup;
            avg_solve += records[i].solve;
            avg_tot   += records[i].setup + records[i].solve;
        }
        avg_setup /= TEST_CNT;
        avg_solve /= TEST_CNT;
        avg_tot   /= TEST_CNT;
        printf("\033[1;35mBEST run:\033[0m Setup time, Solve time, #Iter, Total time\n%f %f %d %f\n",
                        records[best].setup, records[best].solve, records[best].iter, best_tot);
        printf("\033[1;35mAverage run:\033[0m Setup time, Solve time, #Iter, Total time\n%f %f %d %f\n",
                        avg_setup, avg_solve, records[best].iter, avg_tot);
        printf("\033[1;35mFirst run:\033[0m Setup time, Solve time, #Iter, Total time\n%f %f %d %f\n",
                        records[0].setup, records[0].solve, records[0].iter, records[0].setup + records[0].solve);
    }

    free(dist_b); free(dist_x);
    free(dist_row_ptr); free(dist_col_idx); free(dist_vals);

    MPI_Finalize();

    return 0;
}
